package com.xinan.ollama.api.component;

import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.extension.api.R;
import com.xinan.ollama.api.entity.OllamaRequstObj;
import com.xinan.ollama.api.entity.PromptObj;
import com.xinan.ollama.api.utils.OllamaUtil;
import com.xinan.ollama.api.vo.Message;
import com.xinan.ollama.api.vo.req.OllamaReq;
import com.xinan.ollama.api.vo.res.ChatRes;
import com.xinan.ollama.api.vo.res.EmbeddingRes;
import com.xinan.ollama.api.vo.res.GenerateRes;
import com.xinan.ollama.config.OllamaConfig;
import com.xinan.ollama.milvus.MilvusRes;
import com.xinan.ollama.milvus.MilvusUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 问答系统接口
 *
 * @author yuanz
 * @date 2024-08-24
 */
@Slf4j
@Component
public class OllamaComponent {

    @Resource
    private OllamaUtil ollamaUtil;
    @Resource
    private OllamaConfig ollamaConfig;
    @Resource
    private MilvusUtil milvusUtil;

    public R<Object> generate(OllamaReq req) {
        GenerateRes generateRes = null;
        try {
            generateRes = ollamaUtil.generate(req);
        } catch (Exception e) {
            e.printStackTrace();
            log.error("generate Failed!");
        }
        return R.ok(generateRes);
    }

    public R<Object> chat(OllamaReq req) {
        OllamaRequstObj ollamaRequstObj = new OllamaRequstObj();
        ollamaRequstObj.setData(req.getData());
        String reqModel = req.getModel();
        if (StringUtils.isBlank(reqModel)) {
            ollamaRequstObj.setModel(ollamaConfig.model);
        } else {
            ollamaRequstObj.setModel(reqModel);
        }
        ollamaRequstObj.setMessages(req.getMessages());
        ChatRes chatRes = null;
        try {
            chatRes = ollamaUtil.chat(ollamaRequstObj);
        } catch (Exception e) {
            e.printStackTrace();
            log.error("chat Failed!");
        }
        return R.ok(chatRes);
    }

    public R<Object> chatOfRag(OllamaReq req) {
        // 1.请求参数整理
        OllamaRequstObj ollamaRequstObj = new OllamaRequstObj();
        ollamaRequstObj.setData(req.getData());
        String reqModel = req.getModel();
        if (StringUtils.isBlank(reqModel)) {
            ollamaRequstObj.setModel(ollamaConfig.model);
        } else {
            ollamaRequstObj.setModel(reqModel);
        }
        // 2.Prompt设置（RAG）
        List<Float> embedding = getEmbeddingByStr(req.getData());
        List<MilvusRes> milvusRes = milvusUtil.searchContentByFeatureVector("JSGX", Collections.singletonList("JZ"), embedding, 5, 0.5);
        List<Message> messages = new ArrayList<>(2);
        Message system = new Message();
        system.setRole("system");
        final String[] promptContent = {"请根据提供的以下信息：【"};
        milvusRes.forEach(res -> {
            promptContent[0] = promptContent[0] + res.content;
        });
        system.setContent(promptContent[0] + "】回答问题。");
        messages.add(system);
        // 3.用户问题
        Message user = new Message("user", req.getData());
        messages.add(user);
        ollamaRequstObj.setMessages(messages);
        // 4.获取响应
        ChatRes chatRes = null;
        try {
            chatRes = ollamaUtil.chat(ollamaRequstObj);
        } catch (Exception e) {
            e.printStackTrace();
            log.error("chat Failed!");
        }
        return R.ok(chatRes);
    }

    public R<Object> getEventAndOrganizeNameByInformationName(OllamaReq req) {
        return queryOllamaApi(req, PromptObj.PROMPT_OF_EVENT_AND_ORGANIZE_NAME, true);
    }

    public R<Object> chatForKeyPerson(OllamaReq req) {
        return queryOllamaApi(req, PromptObj.PROMPT_OF_KEY_PERSON_V_003, true);
    }

    public R<Object> chatForAnalysisReport(OllamaReq req) {
        return queryOllamaApi(req, PromptObj.PROMPT_OF_ANALYSIS_REPORT, false);
    }

    public R<Object> chatForPersonLevel(OllamaReq req) {
        return queryOllamaApi(req, PromptObj.PROMPT_OF_KEY_PERSON_LEVEL, true);
    }

    public R<Object> generateInformationContentByData(OllamaReq req) {
        return queryOllamaApi(req, PromptObj.PROMPT_OF_GENERATE_INFORMATION_BY_DATA, true);
    }

    public R<Object> generateSubjectAnalysisReportByParam(Map req) {
        Map dataMap = JSON.parseObject(MapUtils.getString(req, "data"),Map.class);;
        List<Map> informationMapList = JSON.parseArray(MapUtils.getString(dataMap, "informationList"), Map.class);
        String informationData = "";
        try {
            informationData = informationMapList.stream().map(m -> MapUtils.getObject(m, "content").toString()).collect(Collectors.joining(";"));
        } catch (Exception e) {
            e.printStackTrace();
        }
        OllamaReq ollamaReq = new OllamaReq();
        ollamaReq.setData(informationData);
        return queryOllamaApi(ollamaReq, PromptObj.PROMPT_OF_ANALYSIS_REPORT, false);
    }

    public R<Object> getEmbeddingByStr(OllamaReq req) {
        EmbeddingRes embeddingRes = null;
        String model = req.getModel();
        if (StringUtils.isBlank(model)) {
            req.setModel("nomic-embed-text:latest");
        }
        try {
            embeddingRes = ollamaUtil.getEmbeddingByStr(req);
        } catch (Exception e) {
            e.printStackTrace();
            log.error("getEmbeddingByStr Failed!");
        }
        return R.ok(embeddingRes);
    }

    public List<Float> getEmbeddingByStr(String str) {
        EmbeddingRes embeddingRes = null;
        OllamaReq req = new OllamaReq();
        req.setModel("nomic-embed-text:latest");
        req.setPrompt(str);
        try {
            embeddingRes = ollamaUtil.getEmbeddingByStr(req);
        } catch (Exception e) {
            e.printStackTrace();
            log.error("getEmbeddingByStr Failed!");
        }
        assert embeddingRes != null;
        return embeddingRes.getEmbedding();
    }


    public R<Object> saveStrEmbeddingData(OllamaReq req) {
        EmbeddingRes embeddingRes = null;
        String model = req.getModel();
        if (StringUtils.isBlank(model)) {
            req.setModel("nomic-embed-text:latest");
        }
        try {
            // 1.获取嵌入向量
            embeddingRes = ollamaUtil.getEmbeddingByStr(req);
            // 2.保存向量数据
            List<Float> embedding = embeddingRes.getEmbedding();
            if (CollectionUtils.isNotEmpty(embedding)) {
                milvusUtil.insertEmbeddingToMilvus("JSGX", "JZ", req.getPrompt(), embedding);
            }
        } catch (Exception e) {
            e.printStackTrace();
            log.error("getEmbeddingByStr Failed!");
        }
        return R.ok(embeddingRes);
    }

    private R<Object> queryOllamaApi(OllamaReq req, String prompt, boolean isReplace) {
        // 0.请求参数整理
        OllamaRequstObj ollamaRequstObj = new OllamaRequstObj();
        List<Message> messages = new ArrayList<>(2);
        // 请求数据
        String reqData = req.getData();
        if (StringUtils.isBlank(reqData)) {
            log.error("reqData is must not null!");
            return R.failed("reqData is must not null!");
        } else {
            ollamaRequstObj.setData(req.getData());
        }
        // 模型设置
        String reqModel = req.getModel();
        if (StringUtils.isBlank(reqModel)) {
            ollamaRequstObj.setModel("qwen2.5:0.5b");
        } else {
            ollamaRequstObj.setModel(reqModel);
        }
        long start = System.currentTimeMillis();
        // Prompt设置
        String reqPrompt = req.getPrompt();
        Message system = new Message();
        system.setRole("system");
        if (StringUtils.isBlank(reqPrompt)) {
            system.setContent(prompt);
        } else {
            system.setContent(reqPrompt);
        }
        messages.add(system);
        // 4.设置内容
        Message user = new Message("user", reqData);
        messages.add(user);
        ollamaRequstObj.setMessages(messages);
        ChatRes chatRes = null;
        try {
            chatRes = ollamaUtil.chat(ollamaRequstObj);
        } catch (Exception e) {
            e.printStackTrace();
            log.error("chat Failed!");
        }
        Message message = chatRes.getMessage();
        String s = "";
        if (message != null) {
            String content = message.getContent();
            String think = "<think>(.|\n)*?</think>";
            s = content.replaceAll(think, "").replace("```json", "").replace("```", "");
            if (isReplace) {
                s = s.replace("\n", "");
            }
            log.info("[{}:{}s]提取信息:{}", ollamaRequstObj.getModel(), (System.currentTimeMillis() - start) / 1000, s);
            return R.ok(s);
        }
        return R.failed("获取数据异常");
    }
}
